import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import shap
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import train_test_split, RandomizedSearchCV


def ensure_picture_folder_exists():
    """Ensure the picture folder exists."""
    picture_folder = "./picture"
    if not os.path.exists(picture_folder):
        os.makedirs(picture_folder)
    return picture_folder


def load_data(file_path):
    """Load data from CSV file."""
    return pd.read_csv(file_path)


def clean_column_names(df):
    """Clean column names by removing spaces."""
    df.columns = df.columns.str.strip()
    return df


def process_medal_column(df):
    """Create 'Gold', 'Silver', 'Bronze', 'Total' columns based on the 'Medal' column."""
    df['Gold'] = df['Medal'].apply(lambda x: 1 if x == 'Gold' else 0)
    df['Silver'] = df['Medal'].apply(lambda x: 1 if x == 'Silver' else 0)
    df['Bronze'] = df['Medal'].apply(lambda x: 1 if x == 'Bronze' else 0)
    df['Total'] = df['Gold'] + df['Silver'] + df['Bronze']
    return df


def convert_noc_to_numeric(df):
    """Convert country codes (NOC) to numeric features."""
    df['NOC'] = df['NOC'].astype('category').cat.codes
    return df


def add_coach_info(df):
    """Add synthetic coach information to the dataset."""
    np.random.seed(42)
    df['Coach'] = np.random.choice(['Coach_A', 'Coach_B', 'Coach_C', 'Coach_D'], size=len(df))
    df['Coach_Experience'] = np.random.randint(1, 20, size=len(df))
    df['Coach_Performance'] = np.random.uniform(0, 1, size=len(df))
    return df


def extract_features(df):
    """Extract features and target variables from the dataset."""
    features = ['NOC', 'Coach_Experience', 'Coach_Performance']
    X = df[features]
    y_gold = df['Gold']
    y_total = df['Total']
    return X, y_gold, y_total


def split_data(X, y_gold, y_total):
    """Split the data into training and testing sets."""
    X_train, X_test, y_train_gold, y_test_gold = train_test_split(X, y_gold, test_size=0.8, random_state=42)
    X_train, X_test, y_train_total, y_test_total = train_test_split(X, y_total, test_size=0.8, random_state=42)
    return X_train, X_test, y_train_gold, y_test_gold, y_train_total, y_test_total


def hyperparameter_tuning(X_train, y_train_gold, y_train_total):
    """Perform hyperparameter tuning using RandomizedSearchCV."""
    param_dist = {
        'n_estimators': [50, 100],
        'learning_rate': [0.05, 0.1],
        'max_depth': [3, 5],
        'min_samples_split': [2, 5],
        'min_samples_leaf': [1, 2],
        'subsample': [0.9, 1.0]
    }

    random_search_gold = RandomizedSearchCV(
        estimator=GradientBoostingRegressor(random_state=42),
        param_distributions=param_dist, n_iter=5, cv=2,
        scoring='neg_mean_squared_error', n_jobs=-1, verbose=1, random_state=42
    )

    random_search_total = RandomizedSearchCV(
        estimator=GradientBoostingRegressor(random_state=42),
        param_distributions=param_dist, n_iter=5, cv=2,
        scoring='neg_mean_squared_error', n_jobs=-1, verbose=1, random_state=42
    )

    random_search_gold.fit(X_train, y_train_gold)
    random_search_total.fit(X_train, y_train_total)

    return random_search_gold.best_estimator_, random_search_total.best_estimator_


def shap_analysis(model, X, features, picture_folder, filename, title):
    """Perform SHAP analysis and visualize SHAP values."""
    explainer = shap.TreeExplainer(model)
    shap_values = explainer.shap_values(X)

    plt.figure(figsize=(10, 6))
    shap.summary_plot(shap_values, X, feature_names=features, show=False)
    plt.title(title)
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/{filename}")
    plt.show()


def visualize_comparison(df, picture_folder, filename, title, xlabel, ylabel):
    """Visualize comparison of total medals before and after improvement."""
    plt.figure(figsize=(10, 6))
    ax = df.plot(kind='bar', figsize=(10, 6), color=['red', 'blue'])
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.xticks(range(len(df)), df.index, rotation=0)
    plt.legend(['Before', 'After'])
    plt.grid(axis='y')
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/{filename}")
    plt.show()


if __name__ == "__main__":
    # Ensure the picture folder exists
    picture_folder = ensure_picture_folder_exists()

    # Step 1: Load data
    athlete_data = load_data("summerOly_athletes.csv")

    # Step 2: Clean column names
    athlete_data = clean_column_names(athlete_data)

    # Step 3: Process Medal column
    athlete_data = process_medal_column(athlete_data)

    # Step 4: Convert NOC to numeric
    athlete_data = convert_noc_to_numeric(athlete_data)

    # Step 5: Add synthetic coach information
    athlete_data = add_coach_info(athlete_data)

    # Step 6: Feature extraction
    X, y_gold, y_total = extract_features(athlete_data)

    # Step 7: Split the data into training and testing sets
    X_train, X_test, y_train_gold, y_test_gold, y_train_total, y_test_total = split_data(X, y_gold, y_total)

    # Step 8: Hyperparameter tuning
    best_model_gold, best_model_total = hyperparameter_tuning(X_train, y_train_gold, y_train_total)

    # Step 9: SHAP analysis for gold medals
    X_test_subset = X_test.sample(frac=0.2, random_state=42)
    shap_analysis(best_model_gold, X_test_subset, ['NOC', 'Coach_Experience', 'Coach_Performance'], picture_folder,
                  "2.1.png", 'SHAP Summary Plot for Gold Medal Prediction')

    # Step 10: SHAP analysis for total medals
    shap_analysis(best_model_total, X_test_subset, ['NOC', 'Coach_Experience', 'Coach_Performance'], picture_folder,
                  "2.2.png", 'SHAP Summary Plot for Total Medal Prediction')

    # Step 11: Select countries with the highest predicted gold medals
    country_gold_predictions = pd.DataFrame({'NOC': athlete_data['NOC'], 'Gold_Prediction': best_model_gold.predict(X)})
    top_3_gold_countries = country_gold_predictions.sort_values(by='Gold_Prediction', ascending=False).head(3)
    print("Top 3 countries with highest gold medal predictions:")
    print(top_3_gold_countries)

    # Step 12: Improve the performance of the "great" coach
    athlete_data['Coach_Performance_Improved'] = athlete_data['Coach_Performance'] * 5
    athlete_data['NOC_Improved'] = athlete_data['NOC'] * 5

    # Step 13: Map the improved features back to the original feature names
    athlete_data['NOC'] = athlete_data['NOC_Improved']
    athlete_data['Coach_Performance'] = athlete_data['Coach_Performance_Improved']

    # Step 14: Update the feature matrix and predict improved total medals
    X_improved = athlete_data[['NOC', 'Coach_Experience', 'Coach_Performance']]
    y_total_improved = best_model_total.predict(X_improved)

    # Step 15: Display the improved results
    top_3_total = top_3_gold_countries[['NOC']].copy()
    top_3_total['Predicted_Total_Before'] = best_model_total.predict(X)[:3]
    top_3_total['Predicted_Total_After'] = y_total_improved[:3]
    print("Improved total medal predictions for the top 3 countries:")
    print(top_3_total)

    # Step 16: Visualization: Show the comparison of total medals before and after improvement
    visualize_comparison(top_3_total[['Predicted_Total_Before', 'Predicted_Total_After']], picture_folder, "2.3.png",
                         'Total Medal Prediction: Before and After Coach Performance Improvement', 'Country',
                         'Total Medal Count')

    # Step 17: Estimate the improvement in total medals
    top_3_total['Total_Improvement'] = top_3_total['Predicted_Total_After'] - top_3_total['Predicted_Total_Before']
    print("Estimated improvement in total medals for the top 3 countries:")
    print(top_3_total[['NOC', 'Total_Improvement']])